package com.hzh.flink.core

import org.apache.flink.api.common.functions.RuntimeContext
import org.apache.flink.api.common.state.{ValueState, ValueStateDescriptor}
import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.functions.KeyedProcessFunction
import org.apache.flink.streaming.api.scala._
import org.apache.flink.util.Collector

object Demo14AvgAge {
  def main(args: Array[String]): Unit = {
    //创建环境
    val env: StreamExecutionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment


    /**
     *
     * 实时计算每一个班级的平均年龄
     */

    /**
     * 读取socket的数据
     */
    val lines: DataStream[String] = env.socketTextStream("master", 8888)

    val clazzAndAge: DataStream[(String, Int)] = lines.map(line => {
      val splits: Array[String] = line.split(",")
      val clazz: String = splits(4)
      val age: Int = splits(2).toInt
      (clazz, age)
    })

    //按照班级分组
    val keyByDS: KeyedStream[(String, Int), String] = clazzAndAge.keyBy(_._1)

    /**
     * 计算平均年龄
     * flink的状态可以在任务算子中使用，map，filter，process都可以（Rich）
     *
     */

    val resultDS: DataStream[(String, Double)] = keyByDS.process(new KeyedProcessFunction[String, (String, Int), (String, Double)] {
      override def open(parameters: Configuration): Unit = {
        /**
         * 定义两个状态来保存总人数和总年龄
         */
        val context: RuntimeContext = getRuntimeContext

        //总人数的状态描述对象
        val sumNumDesc = new ValueStateDescriptor[Int]("sumNum", classOf[Int])

        //总年龄的状态描述对象
        val sumAgeDesc = new ValueStateDescriptor[Int]("sumAge", classOf[Int])

        sumNumState = context.getState(sumNumDesc)

        sumAgeState = context.getState(sumAgeDesc)

      }

      //保存总人数的状态
      var sumNumState: ValueState[Int] = _

      //保存总年龄的状态
      var sumAgeState: ValueState[Int] = _

      override def processElement(value: (String, Int),
                                  ctx: KeyedProcessFunction[String, (String, Int), (String, Double)]#Context,
                                  out: Collector[(String, Double)]): Unit = {

        val clazz: String = value._1
        val age: Int = value._2

        //1、获取之前的总人数和总的年龄
        var sumNum: Int = sumNumState.value()
        //人数累加
        sumNum += 1
        //更新状态
        sumNumState.update(sumNum)


        var sumAge: Int = sumAgeState.value()
        //人数累加
        sumAge += age
        //更新状态
        sumAgeState.update(sumAge)

        //计算平均年龄
        val avgAge: Double = sumAge / sumNum.toDouble

        //将数据发送到下游

        out.collect((clazz, avgAge))

      }
    })

    resultDS.print()

    env.execute()



  }

}
